package org.example.myleet.p692;

import java.util.*;

public class Solution {

    static class Trie {
        //length 26
        Trie[] children;
        String word;
        Integer count;
        Integer idx;

        public Trie() {
            children = new Trie[26];
        }
    }

    static class WordCount {
        String word;
        Integer count;

        public WordCount(String word) {
            this.word = word;
            count = 0;
        }

        public WordCount(String word, Integer count) {
            this.word = word;
            this.count = count;
        }
    }

    private Trie root;

    /**
     * 6 ms
     * 哈希表 + 小根堆
     */
    public List<String> topKFrequent(String[] words, int k) {
        Map<String, Integer> wordCounts = new HashMap<>();
        for (String word : words) {
            Integer count = wordCounts.getOrDefault(word, 0);
            wordCounts.put(word, ++count);
        }
        PriorityQueue<WordCount> heap = new PriorityQueue<>(new Comparator<WordCount>() {
            @Override
            public int compare(WordCount o1, WordCount o2) {
                if (!o1.count.equals(o2.count)) {
                    return o2.count - o1.count;
                }
                return o1.word.compareTo(o2.word);
            }
        });
        for (Map.Entry<String, Integer> entry : wordCounts.entrySet()) {
            heap.offer(new WordCount(entry.getKey(), entry.getValue()));
        }
        List<String> topKWords = new ArrayList<>(k);
        for (int i = 0; i < k; ++i) {
            topKWords.add(heap.poll().word);
        }
        return topKWords;
    }

    /**
     * 10 ms
     * version 2
     * 字典树 + 小根堆
     */
//    public List<String> topKFrequent(String[] words, int k) {
//        root = new Trie();
//        int idx = 0;
//        WordCount[] wordCounts = new WordCount[words.length];
//        for (String word : words) {
//            Trie node = root;
//            for (int i = 0; i < word.length(); ++i) {
//                int c = word.charAt(i) - 'a';
//                if (node.children[c] == null) {
//                    node.children[c] = new Trie();
//                }
//                node = node.children[c];
//            }
//            node.word = word;
//            if (node.count == null) {
//                node.count = 0;
//                wordCounts[idx] = new WordCount(word);
//                node.idx = idx;
//                ++idx;
//            }
//            ++wordCounts[node.idx].count;
//            ++node.count;
//        }
//        PriorityQueue<WordCount> heap = new PriorityQueue<>(new Comparator<WordCount>() {
//            @Override
//            public int compare(WordCount o1, WordCount o2) {
//                if (!o1.count.equals(o2.count)) {
//                    return o2.count - o1.count;
//                }
//                return o1.word.compareTo(o2.word);
//            }
//        });
//        for (int i = 0; i < idx; ++i) {
//            heap.offer(wordCounts[i]);
//        }
//        List<String> topKWords = new ArrayList<>(k);
//        for (int i = 0; i < k; ++i) {
//            topKWords.add(heap.poll().word);
//        }
//        return topKWords;
//    }

    /**
     * 11 ms
     * version 1
     * 字典树 + 数组排序
     */
//    public List<String> topKFrequent(String[] words, int k) {
//        root = new Trie();
//        int idx = 0;
//        WordCount[] wordCounts = new WordCount[words.length];
//        for (String word : words) {
//            Trie node = root;
//            for (int i = 0; i < word.length(); ++i) {
//                int c = word.charAt(i) - 'a';
//                if (node.children[c] == null) {
//                    node.children[c] = new Trie();
//                }
//                node = node.children[c];
//            }
//            node.word = word;
//            if (node.count == null) {
//                node.count = 0;
//                wordCounts[idx] = new WordCount(word);
//                node.idx = idx;
//                ++idx;
//            }
//            ++wordCounts[node.idx].count;
//            ++node.count;
//        }
//        wordCounts = Arrays.copyOf(wordCounts, idx);
//        Arrays.sort(wordCounts, new Comparator<WordCount>() {
//            @Override
//            public int compare(WordCount o1, WordCount o2) {
//                if (!o1.count.equals(o2.count)) {
//                    return o2.count - o1.count;
//                }
//                int len = Math.min(o1.word.length(), o2.word.length());
//                for (int i = 0; i < len; ++i) {
//                    char c1 = o1.word.charAt(i);
//                    char c2 = o2.word.charAt(i);
//                    if (c1 != c2) {
//                        return c1 - c2;
//                    }
//                }
//                return o1.word.length() - o2.word.length();
//            }
//        });
//        List<String> topKWords = new ArrayList<>(k);
//        for (int i = 0; i < k; ++i) {
//            topKWords.add(wordCounts[i].word);
//        }
//        return topKWords;
//    }
}
